# 89. 路径之谜
# https://www.lanqiao.cn/problems/89/learning/?page=1&first_category_id=1&second_category_id=3&contain_answer=true
# Date: 2024/12/31
from leetcode.test import test_function as tf


def back_tracking(n, north, west):
    """我的方法, 使用回溯法, 但是测试用例中只有第一个在规定时间内完成了, 其余的均超时
    优化方面主要考虑剪枝操作, 剪除明显无解的路径
    """
    visited = [[False] * n for _ in range(n)]
    # 设置起点状态
    north[0] -= 1
    west[0] -= 1
    visited[0][0] = True

    def neighbors(r, c) -> (int, int):
        for nr, nc in ((r - 1, c), (r, c - 1), (r + 1, c), (r, c + 1)):
            if 0 <= nr < n and 0 <= nc < n:
                yield nr, nc

    def dfs(path, x: int, y: int):
        nonlocal res, visited
        # 判断是否为答案
        if x == n - 1 and y == n - 1:
            if all(num == 0 for num in north) and all(num == 0 for num in west):
                res = path[:]
                return
        for r, c in neighbors(x, y):
            if visited[r][c]:
                continue
            # 进入一个新节点的操作
            path.append(n * r + c)
            visited[r][c] = True
            north[c] -= 1
            west[r] -= 1
            # 遍历后面的节点
            dfs(path, r, c)
            # 回溯操作
            path.pop()
            visited[r][c] = False
            north[c] += 1
            west[r] += 1

    res = []
    dfs([0], 0, 0)
    return res


def back_tracking_opt(n, north, west):
    """优化的方法
    1. 添加提前退出的条件
    2. 添加剪枝的过程: 在路径中判断是否该路径已经错误
    3. 将visited二维数组修改为一维数组以减少检索的时间
    """
    # 优化:修改为一维数组
    visited = [False] * (n * n)
    # 设置起点状态
    north[0] -= 1
    west[0] -= 1
    visited[0] = True

    def neighbors(r, c) -> (int, int):
        for nr, nc in ((r - 1, c), (r, c - 1), (r + 1, c), (r, c + 1)):
            if 0 <= nr < n and 0 <= nc < n:
                yield nr, nc

    def dfs(path, x: int, y: int):
        nonlocal res, found
        # # 优化: 已经找到答案, 提前返回
        if found:
            return

        # 判断是否为答案
        if x == n - 1 and y == n - 1:
            if all(num == 0 for num in north) and all(num == 0 for num in west):
                res = path[:]
                found = True
                return
        for nx, ny in neighbors(x, y):
            idx = n * nx + ny  # 转换为一维数组的索引
            if visited[idx]:
                continue

            # 优化: 剪枝, 排除一些必然错误的选项
            if north[ny] <= 0 or west[nx] <= 0:
                continue

            # 进入一个新节点的操作
            path.append(idx)
            visited[idx] = True
            north[ny] -= 1
            west[nx] -= 1

            # 递归搜索
            dfs(path, nx, ny)

            # 回溯操作
            path.pop()
            visited[idx] = False
            north[ny] += 1
            west[nx] += 1

    res = []
    found = False
    dfs([0], 0, 0)
    return res


def back_tracking_ultra(n, north, west):
    """一种优化的剪枝方法
    在第i个靶子已经为0的情况下, 若i前面仍有靶子不为0, 那么就一定是一个错误的路径, 因为目的地在右下角, 而如果此前仍有非零的数字,
    那么就必然要向左(或上)移动, 但是还是要移动回来, 就必须要经过该列(或行), 那么这个为此时为0的列(或行)在未来就一定小于0, 不符合题意.
    如:
    1 2 0 3
      <-x
    此时为了满足第1/2列为0的条件, x需要向左移动:
    0 0 0 3
      x->
    满足了1/2列之后, x若继续向右移动, 则不满足题意

    虽然上面的优化思路很好, 但是由于判断过于复杂, 在某些情况下甚至不如此前的剪枝方法, 所以只作为一种思路记录
    """
    visited = [False] * (n * n)
    # 设置起点状态
    north[0] -= 1
    west[0] -= 1
    visited[0] = True

    def neighbors(r, c) -> (int, int):
        for nr, nc in ((r - 1, c), (r, c - 1), (r + 1, c), (r, c + 1)):
            if 0 <= nr < n and 0 <= nc < n:
                yield nr, nc

    def dfs(path, x: int, y: int):
        nonlocal res, found
        # 优化: 已经找到答案, 提前返回
        if found:
            return

        # 判断是否为答案
        if x == n - 1 and y == n - 1:
            if all(num == 0 for num in north) and all(num == 0 for num in west):
                res = path[:]
                found = True
                return

        for nx, ny in neighbors(x, y):
            idx = n * nx + ny  # 转换为一维数组的索引
            if visited[idx]:
                continue

            # 优化: 剪枝, 排除一些必然错误的选项
            if north[ny] <= 0 or west[nx] <= 0:
                continue
            if north[ny] == 0 and any(north[i] > 0 for i in range(ny)):
                continue
            if west[nx] == 0 and any(west[i] > 0 for i in range(nx)):
                continue

            # 进入一个新节点的操作
            path.append(idx)
            visited[idx] = True
            north[ny] -= 1
            west[nx] -= 1

            # 递归搜索
            dfs(path, nx, ny)

            # 回溯操作
            path.pop()
            visited[idx] = False
            north[ny] += 1
            west[nx] += 1

    res = []
    found = False
    dfs([0], 0, 0)
    return res


if __name__ == '__main__':
    """使用回溯法"""
    # n = int(input())
    # north = [int(i) for i in input().split()]
    # west = [int(i) for i in input().split()]
    n = 4
    north = [2, 4, 3, 4]
    west = [4, 3, 3, 3]
    ans = [0, 4, 5, 1, 2, 3, 7, 11, 10, 9, 13, 14, 15]

    for item in back_tracking_opt(n, north, west):
        print(item, end=' ')

    tf(back_tracking, input_args={"n": n, "north": north, "west": west},
       expected_results=ans, times=1000)
    tf(back_tracking_opt, input_args={"n": n, "north": north, "west": west},
       expected_results=ans, times=1000)
    tf(back_tracking_ultra, input_args={"n": n, "north": north, "west": west},
       expected_results=ans, times=1000)

    n = 5
    north = [3, 2, 3, 4, 1]
    west = [1, 4, 3, 2, 3]
    ans = [0, 5, 10, 11, 6, 7, 8, 13, 18, 17, 22, 23, 24]
    tf(back_tracking_opt, input_args={"n": n, "north": north, "west": west},
       expected_results=ans, times=1000)
    tf(back_tracking_ultra, input_args={"n": n, "north": north, "west": west},
       expected_results=ans, times=1000)
